'''
author: zpj
图像处理类

类的用法说明
ImageProcessing 类提供了一系列图像处理和 OCR 识别相关的功能，主要包括截图、图像灰度化、OCR 识别以及页面切换等操作。下面是该类的使用步骤和示例：
步骤
初始化类：创建 ImageProcessing 类的实例，需要传入页面名称作为参数，该名称用于生成截图和灰度图的文件名。
截图：调用 screenshot 方法进行全屏截图，并将截图保存到指定路径。
图像灰度化：调用 grayscale 方法将截图转换为灰度图，并保存到指定路径。
OCR 识别：调用 ocr_recognition 方法对指定图像进行 OCR 识别，默认使用灰度图。
直接截图识别：调用 ocr_recognition_screenshot 方法进行全屏截图并直接进行 OCR 识别，不保存截图。
页面切换：调用 switch_page 方法，传入目标页面名称，程序会尝试在当前屏幕中找到该页面名称并点击其中心位置。

示例

----------- 分步示例 会保存截图、灰度图、OCR 识别结果到 images 文件夹中 -----------------
# 初始化 ImageProcessing 类
image_processing = ImageProcessing("build_dino")

# 截图
image_processing.screenshot()

# 灰度化
image_processing.grayscale()

# OCR 识别
result = image_processing.ocr_recognition()
print(result)
for line in result[0]:
    print(line)
----------------------------

# 直接截取当前屏幕并进行OCR识别 (不保存截图) 
direct_result = ImageProcessing.ocr_recognition_screenshot()
print(direct_result)

# 程序会尝试在当前屏幕中找到该页面名称并点击其中心位置。
ImageProcessing.switch_page("build_dino")

'''
import os
import cv2
import numpy as np
from paddleocr import PaddleOCR
from PIL import ImageGrab  # 用于截图
import pyautogui  # 用于鼠标点击
import time  # 用于延时

class ImageProcessing:

    def __init__(self, page_name):
        '''
        初始化
        :param page_name: 页面名称，用于生成截图和灰度图的文件名
        '''
        self.page_name = page_name
        self.image_path = f"images/{page_name}.png"
        self.image_path_gray = f"images/{page_name}_gray.png"

    # 截图
    def screenshot(self):
        '''
        全屏截图
        '''
        # 查看images文件夹是否存在，不存在则创建
        if not os.path.exists("images"):
            os.makedirs("images")

        # 查看是否已经存在截图
        if os.path.exists(self.image_path):
            print(f"已存在截图: {self.image_path}")
            return
        
        screenshot = ImageGrab.grab()  # 全屏截图
        screenshot = np.array(screenshot)  # 转换为 NumPy 数组
        screenshot = cv2.cvtColor(screenshot, cv2.COLOR_RGB2BGR)  # 转换为 BGR 格式

        # 保存截图
        cv2.imwrite(self.image_path, screenshot)
        print(f"截图已保存到: {self.image_path}")

    # 图像灰度化
    def grayscale(self):
        '''
        图像灰度化
        :return: 返回灰度化后的图像
        '''
        # 灰度图是否存在
        if os.path.exists(self.image_path_gray):
            print(f"已存在灰度图: {self.image_path_gray}")
            return

        # 读取截图
        if not os.path.exists(self.image_path):
            raise FileNotFoundError(f"未找到截图文件: {self.image_path}")

        image = cv2.imread(self.image_path)  # 读取截图
        if image is None:
            raise ValueError(f"无法读取图片: {self.image_path}")

        # 转换为灰度图
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 保存灰度图
        cv2.imwrite(self.image_path_gray, gray_image)
        print(f"灰度图已保存到: {self.image_path_gray}")

    # OCR 识别
    def ocr_recognition(self, img_path=None):
        '''
        OCR 识别
        :param img_path: 输入的图像路径（可选），如果为 None，则使用灰度图
        :return: OCR 识别结果
        '''
        if img_path is None:
            img_path = self.image_path_gray  # 默认使用灰度图

        if not os.path.exists(img_path):
            raise FileNotFoundError(f"未找到图片文件: {img_path}")

        # 读取图片
        img = cv2.imread(img_path)
        if img is None:
            raise ValueError(f"无法读取图片: {img_path}")

        ocr = PaddleOCR(use_angle_cls=True, lang='ch')  # 初始化 PaddleOCR
        result = ocr.ocr(img, cls=True)  # 进行 OCR 识别
        return result
    
    # 截图不保存直接识别ocr
    @staticmethod
    def ocr_recognition_screenshot():
        '''截图不保存直接识别ocr
        '''
        screenshot = ImageGrab.grab()  # 全屏截图
        screenshot = np.array(screenshot)  # 转换为 NumPy 数组
        screenshot = cv2.cvtColor(screenshot, cv2.COLOR_RGB2BGR)  # 转换为 BGR 格式

        # 灰度化
        gray_image = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)

        # OCR 识别
        ocr = PaddleOCR(use_angle_cls=True, lang='ch')  # 初始化 PaddleOCR
        result = ocr.ocr(gray_image, cls=True)  # 进行 OCR 识别

        print(result)
        return result[0]
    

    # 切换到指定页面
    @staticmethod
    def switch_page(page_name):
        '''截取当前屏幕并进行OCR识别，找到指定页面的中心位置并点击'''

        ocr_result = ImageProcessing.ocr_recognition_screenshot()

        for i in ocr_result:
            if i[1][0] == page_name:
                # 提取四个角的坐标点
                ocr_coordinates = i[0]

                # 计算中心坐标点
                x_coords = [point[0] for point in ocr_coordinates]  # 提取所有 x 坐标
                y_coords = [point[1] for point in ocr_coordinates]  # 提取所有 y 坐标

                center_x = sum(x_coords) / 4  # 计算 x 中心点
                center_y = sum(y_coords) / 4  # 计算 y 中心点

                # 点击坐标点
                pyautogui.click(x=center_x, y=center_y)
                time.sleep(0.5)
                break

if __name__ == '__main__':
    # 测试
    image_processing = ImageProcessing("build_dino")

    # 截图
    image_processing.screenshot()

    # 灰度化
    image_processing.grayscale()

    # OCR 识别
    result = image_processing.ocr_recognition()
    print(result)
    for line in result[0]:
        print(line)